#include "FramePrediction.h"
#include "CoreGlobals.h"
#include "ImageUtils.h"
#include "GameFramework/GameUserSettings.h"
#include "Kismet/GameplayStatics.h"
#include "Camera/CameraComponent.h"
#include "Engine/Engine.h"
#include "OpenGLDrv.h"
#include "OpenGLResources.h"
#include "PostProcess/SceneRenderTargets.h"

#if PLATFORM_ANDROID
#define FF_DECLARE_PROTOTYPES 1
#include <GLES3/gl31.h>
#include "frame_flow_core.h"
#include "frame_flow_gl.h"
static FFInstance s_Instance = nullptr;
#endif

bool FramePrediction::bRealFrame = true;
bool FramePrediction::bColorBufferInit = false;
bool FramePrediction::bEnable = false;  // default disable
bool FramePrediction::bNativeFFInit = false;

DEFINE_LOG_CATEGORY_STATIC(LogFFInterpolation, Verbose, All);

GLint g_FP_drawFboId = 0;  // Global Variables, used for GL state backup and recovery
GLint g_FP_readFboId = 0;

TAutoConsoleVariable<int32> CVarEnableFramePredictionFP(
	TEXT("r.FramePrediction.FP.Enable"),
	0,
	TEXT("Switch for Frame Prediction"),
	ECVF_RenderThreadSafe);

void FramePrediction::UpdateEnableFromSetting()
{
#if PLATFORM_ANDROID
	bool currEnable = CVarEnableFramePredictionFP.GetValueOnAnyThread() > 0;
	if (FramePrediction::bEnable != currEnable) {
		FramePrediction::bEnable = currEnable;
		UE_LOG(LogScript, Warning, TEXT("FramePrediction Change Enable: %d"), bEnable);
	}
#endif
}

FramePrediction::FramePrediction() : FFDataManager()
{
    FCoreDelegates::OnPostEngineInit.AddRaw(this, &FramePrediction::OnPostEngineInit);
}

FramePrediction::~FramePrediction()
{

}

void FramePrediction::OnSetEnable(bool Enable)
{
	UE_LOG(LogScript, Warning, TEXT("FramePrediction OnSetEnable: %d"), Enable);
	bEnable = Enable;
}

void FramePrediction::InitFF()
{
#if PLATFORM_ANDROID
	// Create FF instance
	FFInstanceCreateInfo createInfo
	{
		.sType = FF_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
		.pNext = nullptr,
		.majorVersion = 5U,
		.minorVersion = 0U,
		.pfnApiLoaderFunction = nullptr,
		.pfnErrorReportCallback = nullptr,
		.pUserData = nullptr
	};

	s_Instance = FrameFlow_Create(&createInfo);
	if (!s_Instance)
	{
		UE_LOG(LogFFInterpolation, Log, TEXT("FFrameFlowInterpolation s_Instance failed"));
		return;
	}

	// Init FF instance
	FFPredictionAlgorithmInfo algorithmInfo{
		.sType = FF_STRUCTURE_TYPE_PREDICTION_ALGORITHM_INFO,
		.pNext = nullptr,
		.algorithm = FF_PREDICTION_ALGORITHM_HYDRA,
		.extrapolation = FF_FALSE
	};
	FFResult result = FrameFlow_SetPredictionAlgorithm(s_Instance, &algorithmInfo, nullptr);
	UE_LOG(LogFFInterpolation, Log, TEXT("FFrameFlowInterpolation FrameFlow_SetPredictionAlgorithm %d"), result);

	FIntPoint Resolution = GEngine->GetGameUserSettings()->GetScreenResolution();
	UE_LOG(LogFFInterpolation, Log, TEXT("FFrameFlowInterpolation Resolution %d %d"), Resolution.X, Resolution.Y);
	FFExtent2D const resolution{
		.width = static_cast<uint32_t>(Resolution.X),
		.height = static_cast<uint32_t>(Resolution.Y)
	};
	result = FrameFlow_SetResolution(s_Instance, &resolution);
	UE_LOG(LogFFInterpolation, Log, TEXT("FFrameFlowInterpolation FrameFlow_SetResolution %d"), result);

	result = FrameFlow_SetCvvZSemantic(s_Instance, FF_CVV_Z_SEMANTIC_MINUS_ONE_TO_ONE_REVERSE_Z);
	UE_LOG(LogFFInterpolation, Log, TEXT("FFrameFlowInterpolation FrameFlow_SetCvvZSemantic %d"), result);

	result = FrameFlow_Activate(s_Instance);
	UE_LOG(LogFFInterpolation, Log, TEXT("FFrameFlowInterpolation FrameFlow_Activate %d"), result);
#endif
}

void FramePrediction::OnPostEngineInit()
{
	UE_LOG(LogScript, Warning, TEXT("OnPostEngineInit"));
	if (FSlateApplication::IsInitialized())
	{
		FSlateApplication& App = FSlateApplication::Get();
		FSlateRenderer* SlateRenderer = App.GetRenderer();
#if PLATFORM_ANDROID
		SlateRenderer->OnSlateWindowRendered().AddRaw(const_cast<FramePrediction*>(this), &FramePrediction::OnSlateWindowRendered);
		GEngine->GetPostBasePassDelegate().AddRaw(const_cast<FramePrediction*>(this), &FramePrediction::InterpolateFrameEx);
		GEngine->GetPostPostProcessingDelegate().AddRaw(const_cast<FramePrediction*>(this), &FramePrediction::OnPostPostProcessingCallback);
		UE_LOG(LogScript, Warning, TEXT("OnPostEngineInit at Android platform"));
#endif
	}
}

void FramePrediction::OnSlateWindowRendered(SWindow& SlateWindow, void* ViewportRHIPtr)
{
	UpdateEnableFromSetting();
	if (!bEnable) {
		return;
	}

	if (!bColorBufferInit) {
		UE_LOG(LogScript, Warning, TEXT("OnSlateWindowRendered, color buffer not init, return."));
		return;
	}

	if (!IsInGameThread()) {
		UE_LOG(LogScript, Warning, TEXT("OnSlateWindowRendered Not in game thread, return."));
		return;
	}

	if (!bRealFrame) {
		bRealFrame = true;
		FViewportRHIRef Viewport = *((FViewportRHIRef*)ViewportRHIPtr);
		FSlateApplication& App = FSlateApplication::Get();
		TSharedPtr<SWindow> WindowPtr;
		TSharedPtr<SWidget> TestWidget = SlateWindow.AsShared();
		while (TestWidget && !WindowPtr.IsValid()) {
			if (TestWidget->Advanced_IsWindow()) {
				WindowPtr = StaticCastSharedPtr<SWindow>(TestWidget);
			}
			TestWidget = TestWidget->GetParentWidget();
		}

		// Call native FF
		FrameFlowDrawFrame();

		// Copy FF result to Backbuffer
		ENQUEUE_RENDER_COMMAND(OnSlateWindowRendered)([this, Viewport](FRHICommandListImmediate& RHICmdList)
		{
			FTexture2DRHIRef BackBuffer = RHICmdList.GetViewportBackBuffer(Viewport);
			FResolveParams ResolveParams;
			ResolveParams.Rect = FResolveRect(0, 0, BackBuffer->GetSizeX(), BackBuffer->GetSizeY());
			ResolveParams.DestRect = FResolveRect(0, BackBuffer->GetSizeY(), BackBuffer->GetSizeX(), 0);
			if (!FFOutputRT) {
				UE_LOG(LogScript, Warning, TEXT("OnSlateWindowRendered FFOutputRT is null"));
				return;
			}
			CopyTexture(FFOutputRT->GetTargetableRHI(), BackBuffer, false, ResolveParams);
		});

		App.ForceRedrawWindow(WindowPtr.ToSharedRef());  // very important, reduce CPU load
	}
	else {
		bRealFrame = false;
	}
}

void FramePrediction::FrameFlowDrawFrame()
{
#if PLATFORM_ANDROID
	ENQUEUE_RENDER_COMMAND(FrameFlowDrawFrame)([this](FRHICommandListImmediate& RHICmdList)
	{
		if (!FFDataManager.GetCurrent() || !FFDepthRT)
		{
			UE_LOG(LogScript, Warning, TEXT("FrameFlowDrawFrame Color buffer or depth buffer is null"));
			return;
		}

		uint32_t InputColorId = *reinterpret_cast<uint32_t*>(FFDataManager.GetCurrent()->GetTargetableRHI()->GetNativeResource());
		uint32_t OutputColorId = *reinterpret_cast<uint32_t*>(FFOutputRT->GetTargetableRHI()->GetNativeResource());
		uint32_t DepthId = *reinterpret_cast<uint32_t*>(FFDepthRT->GetTargetableRHI()->GetNativeResource());

		FFPerFrameData perFrameData_{
			.sType = FF_STRUCTURE_TYPE_PER_FRAME_DATA,
			.pNext = nullptr,

			.flags = FF_PER_FRAME_DATA_DEPTH_BIT | FF_PER_FRAME_DATA_SCENE_COLOR_BIT |
				FF_PER_FRAME_DATA_VIEW_PROJ_BIT | FF_PER_FRAME_DATA_INV_VIEW_PROJ_BIT,

			.sceneColor = static_cast<uint32_t>(InputColorId),
			.depthStencil = static_cast<uint32_t>(DepthId),
			.viewProj = {},
			.invViewProj = {}
		};

		for (int32 i = 0; i < 4; i++) {
			for (int32 j = 0; j < 4; j++) {
				perFrameData_.viewProj.data[4 * i + j] = ViewProjectionMatrix.M[i][j];
				perFrameData_.invViewProj.data[4 * i + j] = InvViewProjectionMatrix.M[i][j];
			}
		}

		FFDrawFrameInfo drawFrameInfo_{
			.sType = FF_STRUCTURE_TYPE_DRAW_FRAME_INFO,
			.pNext = nullptr,
			.target = OutputColorId,
			.targetIsTexture = FF_TRUE,
			.doDraw = FF_TRUE
		};

		RHICmdList.EnqueueLambda([this, perFrameData_, drawFrameInfo_](FRHICommandListImmediate& cmd)
		{
			FFResult result = FrameFlow_ProvidePerFrameData(s_Instance, &perFrameData_);
			UE_LOG(LogFFInterpolation, Log, TEXT("FFrameFlowInterpolation FrameFlow_ProvidePerFrameData %d"), result);
			result = FrameFlow_DrawFrame(s_Instance, &drawFrameInfo_);
			UE_LOG(LogFFInterpolation, Log, TEXT("FFrameFlowInterpolation FrameFlow_DrawFrame %d"), result);
		});
	});
#endif
}

void FramePrediction::CopyTexture(FRHITexture* Source, FRHITexture* Dest, bool IsSetSourceFrameBuffer, FResolveParams ResolveParams)
{
#if PLATFORM_ANDROID
	FRHICommandListImmediate& RHICmdList = FRHICommandListExecutor::GetImmediateCommandList();

	if (IsSetSourceFrameBuffer)
	{
		RHICmdList.EnqueueLambda([](FRHICommandListImmediate& cmd)
		{
			// Backup GL state
			GLint drawFboId = 0;
			GLint readFboId = 0;
			glGetIntegerv(GL_DRAW_FRAMEBUFFER_BINDING, &drawFboId);
			glGetIntegerv(GL_READ_FRAMEBUFFER_BINDING, &readFboId);
			g_FP_drawFboId = drawFboId;
			g_FP_readFboId = readFboId;
		});
	}

	RHICmdList.CopyToResolveTarget(Source, Dest, ResolveParams);

	if (IsSetSourceFrameBuffer)
	{
		RHICmdList.EnqueueLambda([](FRHICommandListImmediate& cmd)
		{
			// Recovery GL state
			GLint drawFboId = g_FP_drawFboId;
			GLint readFboId = g_FP_readFboId;
			glBindFramebuffer(GL_DRAW_FRAMEBUFFER_BINDING, drawFboId);
			glBindFramebuffer(GL_READ_FRAMEBUFFER_BINDING, readFboId);
		});
	}
#endif
}

void FramePrediction::InterpolateFrameEx(FRDGBuilder& GraphBuilder, const FViewInfo& View)
{
	if (!bEnable) {
		return;
	}

	FViewMatrices ViewMatrices = View.ViewMatrices;
	ViewMatrix = ViewMatrices.GetViewMatrix();
	ProjectionMatrix = ViewMatrices.GetProjectionMatrix();
	ViewProjectionMatrix = ViewMatrices.GetViewProjectionMatrix();
	ProjectionMatrix = AdjustProjectionMatrix(ProjectionMatrix);
	ViewProjectionMatrix = AdjustProjectionMatrix(ViewProjectionMatrix);
	InvViewProjectionMatrix = ViewProjectionMatrix.Inverse();
	UE_LOG(LogFFInterpolation, Log, TEXT("FFrameFlowInterpolation Adjust Camera View Matrix: %s, Projection Matrix: %s, ViewProjectionMatrix: %s"),
		*ViewMatrix.ToString(), *ProjectionMatrix.ToString(), *ViewProjectionMatrix.ToString());

	ENQUEUE_RENDER_COMMAND(InterpolateFrameEx)([this](FRHICommandListImmediate& RHICmdList)
	{
		UE_LOG(LogScript, Warning, TEXT("InterpolateFrameEx"));

		// Scene Depth
		FSceneRenderTargets& SceneContext = FSceneRenderTargets::Get(RHICmdList);
		FTexture2DRHIRef SceneDepthRef = SceneContext.GetSceneDepthSurface();

		if (!FFDepthRT && (SceneDepthRef->GetSizeX() != SceneDepthRef->GetSizeY()))
		{
			UE_LOG(LogScript, Warning, TEXT("InterpolateFrameEx Create FFDepthRT"));
			FRHIResourceCreateInfo CreateInfo;
			FString DebugName = "FFDepthRT";
			CreateInfo.DebugName = *DebugName;
			UE_LOG(LogScript, Warning, TEXT("First Depth width: %d, height: %d"), SceneDepthRef->GetSizeX(), SceneDepthRef->GetSizeY());

			FRHICopyTextureInfo Info;
			Info.Size.X = SceneDepthRef->GetSizeX();
			Info.Size.Y = SceneDepthRef->GetSizeY();

			FPooledRenderTargetDesc RTDesc(FPooledRenderTargetDesc::Create2DDesc(FIntPoint(Info.Size.X, Info.Size.Y),
				SceneDepthRef->GetFormat(),
				FClearValueBinding::None,
				TexCreate_None,
				TexCreate_UAV | TexCreate_ShaderResource | TexCreate_DepthStencilTargetable,
				false,
				1,
				true,
				true));

			GRenderTargetPool.FindFreeElement(RHICmdList, RTDesc, FFDepthRT, TEXT("FFDepthRT"));
		}

		if (SceneDepthRef->GetSizeX() == SceneDepthRef->GetSizeY())
		{
			// ��ʱ���bug��ǰ��֡���ȡ�������depth buffer size������������й�
			UE_LOG(LogScript, Warning, TEXT("Get wrong depth buffer size, return."));
			return;
		}

		UE_LOG(LogScript, Warning, TEXT("Depth width: %d, height: %d"), SceneDepthRef->GetSizeX(), SceneDepthRef->GetSizeY());
		FResolveParams ResolveParams;
		ResolveParams.Rect.X1 = 0;
		ResolveParams.Rect.Y1 = 0;
		ResolveParams.Rect.X2 = SceneDepthRef->GetSizeX();
		ResolveParams.Rect.Y2 = SceneDepthRef->GetSizeY();
		ResolveParams.DestRect.X1 = 0;
		ResolveParams.DestRect.Y1 = 0;
		ResolveParams.DestRect.X2 = SceneDepthRef->GetSizeX();
		ResolveParams.DestRect.Y2 = SceneDepthRef->GetSizeY();
		CopyTexture(SceneDepthRef, FFDepthRT->GetTargetableRHI(), true, ResolveParams);
	});
}

void FramePrediction::OnPostPostProcessingCallback(FRDGBuilder& GraphBuilder)
{
	if (!bEnable) {
		return;
	}

	bColorBufferInit = true;

	UWorld* World = GWorld->GetWorld();
	if (!World)
	{
		UE_LOG(LogFFInterpolation, Log, TEXT("FFrameFlowInterpolation Failed to get UWorld."));
		return;
	}

	ENQUEUE_RENDER_COMMAND(OnPostPostProcessingCallback)([this](FRHICommandListImmediate& RHICmdList)
	{
		if (!bNativeFFInit) {
			RHICmdList.EnqueueLambda([this](FRHICommandListImmediate& RHICmdList)
				{
					InitFF();
				});
			bNativeFFInit = true;
		}

		auto* Engine = GEngine;
		auto GameViewport = Engine ? Engine->GameViewport : nullptr;
		auto Viewport = GameViewport ? GameViewport->Viewport : nullptr;
		auto ViewportRHI = Viewport ? Viewport->GetViewportRHI() : nullptr;
		FTexture2DRHIRef BackBuffer = RHIGetViewportBackBuffer(ViewportRHI);
		if (!BackBuffer) {
			UE_LOG(LogFFInterpolation, Log, TEXT("Failed to get backbuffer."));
			return;
		}

		FRHICopyTextureInfo Info;
		Info.Size.X = BackBuffer->GetSizeX();
		Info.Size.Y = BackBuffer->GetSizeY();

		// Create resource for RenderTarget
		FPooledRenderTargetDesc RTDesc(FPooledRenderTargetDesc::Create2DDesc(FIntPoint(Info.Size.X, Info.Size.Y),
			BackBuffer->GetFormat(),
			FClearValueBinding::None,
			TexCreate_None,
			TexCreate_UAV | TexCreate_ShaderResource,
			false,
			1,
			true,
			true));

		for (int32 index = 0; index < FFDataManager.FFInputRT.Num(); index++) {
			FString RTName = "FramePredictionInterpolated" + FString::FromInt(index);
			if (FFDataManager.FFInputRT[index] == nullptr) {
				GRenderTargetPool.FindFreeElement(RHICmdList, RTDesc, FFDataManager.FFInputRT[index], *RTName);
			}
		}
		if (FFOutputRT == nullptr) {
			GRenderTargetPool.FindFreeElement(RHICmdList, RTDesc, FFOutputRT, TEXT("FFOutputRT"));
		}

		FFDataManager.SwapBuffers();
		FResolveParams ResolveParams;
		ResolveParams.Rect = FResolveRect(0, 0, BackBuffer->GetSizeX(), BackBuffer->GetSizeY());
		ResolveParams.DestRect = FResolveRect(0, BackBuffer->GetSizeY(), BackBuffer->GetSizeX(), 0);
		CopyTexture(BackBuffer, FFDataManager.GetCurrent()->GetTargetableRHI(), true, ResolveParams);
		CopyTexture(FFDataManager.GetPrevious()->GetTargetableRHI(), BackBuffer, false, ResolveParams);
	});
}